There is partial overlap with #1006, which addresses one slice of this problem space (0% cache reuse for hybrid models with non-trimmable cache types). Both PRs were developed independently in response to #980.

#1006 handles the non-trimmable checkpoint case. This PR also handles thinking-token divergence (where chat templates strip <think> content between turns), rewind-based longer-cache reuse for sliding-window models, refcounted extraction for concurrent request safety, and batch checkpoint compatibility for mixed-length prefill batches. Test coverage is ~1,800 lines across 6 test files, including end-to-end logprob-equivalence checks against a real model.

Benchmark methodology and results are in the PR description. Harness source is at lyonsno/mlx-lm-benchmarks.

Nice work. We run Qwen3.5-122B with thinking enabled on M2 Ultra 128GB and the 0% cache reuse on multi-turn is one of our top pain points. The checkpoint-at-think-boundary approach is the right call.

Code review notes:

Strengths:

• Test coverage is excellent. The real-model logprob-equivalence tests with Qwen 0.5B give strong correctness confidence. Mock layers for edge cases (partial rewind failure, legacy fallback) cover the important failure modes.
• The BatchRotatingKVCache.trim() fix for left_padding on rotated caches looks like a genuine bug fix that should land regardless of the rest of the PR.
• Fail-closed design on rewind: deepcopy, attempt rewind, discard on failure. Original cache stays intact.

Minor concerns:

• _can_rewind_layer_cache introspects type(layer_cache).rewind is not _BaseCache.rewind to detect custom overrides. This is fragile if someone subclasses and delegates without overriding. A simpler protocol check (hasattr + explicit supports_rewind flag, or just try/except) might be more robust.
• The count field on CacheEntry reads like a refcount but is really an insertion counter. A rename to insertion_count or a brief docstring would help readers.
• insert_cache no longer evicts shorter prefixes. This is correct for checkpoints but increases LRU pressure in non-thinking scenarios. Probably fine in practice since LRU eviction handles it, but worth noting.

MTP compatibility: The changes to cache.py are additive (new methods on BatchRotatingKVCache). Since we use an editable mlx-lm with MTP support (PR #990), these would flow in cleanly. The standard KVCache used by MTP layers has is_trimmable()=True + trim(), so the legacy fallback path handles them correctly.

We can benchmark on the 122B (M2 Ultra 128GB) once this is ready to test. Happy to report TTFT at turn 5/10/20 with thinking mode on Qwen3.5.

Production TTFT data from M2 Ultra 128GB, Qwen3.5-122B-A10B 5-bit, thinking mode enabled.

10-turn coding conversation (progressive refinement of a Python module):

TTFT growth factor: 0.95x (flat, not growing). We're running vllm-mlx's prefix cache, not mlx-lm's LRUPromptCache, so the comparison isn't direct. But it confirms the problem: turns with heavy thinking content (1-2, 6-10) show 25-28s TTFT on 122B. The checkpoint approach in this PR would directly address the thinking-token divergence causing cache misses.

We'll do a proper A/B test with the full PR applied when it's closer to merge. For now the cache.py changes (BatchRotatingKVCache rewind/trim fix) are applied to our editable mlx-lm without issues.

Thanks for the thoughtful review. I incorporated the two code-level suggestions that seemed worth tightening.

For the rewind-path check, I refactored the _BaseCache contract so it now has a default rewind() plus a _has_rewind_impl() helper that uses method identity to detect real overrides (6070f37). That keeps the existing fail-closed behavior while making the intent clearer, and it means third-party _BaseCache subclasses that implement rewind() participate automatically without needing an extra opt-in flag. I also added targeted tests covering the base class, a subclass without an override, a subclass with a real override, and BatchRotatingKVCache.

I also renamed CacheEntry.count to ref_count. You were right that count was too vague; this field is decremented on extraction, so ref_count is a better description.

On shorter-prefix retention, I left the behavior unchanged. That was a deliberate tradeoff rather than an oversight: checkpoint entries need shorter prefixes to coexist, and LRU eviction already bounds the pressure. If someone later shows a workload where coexisting shorter prefixes measurably hurt hit rate or memory behavior, I'd be happy to revisit that separately.

Good to hear that the fallback path through is_trimmable()/trim() lines up cleanly with your MTP setup. And yes, I'd definitely welcome the 122B A/B when you get a chance. TTFT at turns 5/10/20 with thinking enabled would be especially useful.

Merged this PR into our editable mlx-lm (feat/mtp-native branch, Qwen3.5-122B). The merge was clean -- one trivial conflict in qwen3_5.py (our MTP _process_chunk refactor vs your cache.advance(S) addition, both kept).

15-turn benchmark on M2 Ultra 128GB, Qwen3.5-122B-A10B 5-bit, thinking enabled, vllm-mlx BatchedEngine:

TTFT growth factor: 0.64x (TTFT not growing across turns). Note: TTFTs elevated vs earlier run due to concurrent load testing sharing the server.

The ref_count rename and _has_rewind_impl refactor look good. The BatchRotatingKVCache.rewind() and trim() left_padding fix are clean and additive -- no issues with our MTP cache layers.

For a proper isolated A/B (with vs without PR, same workload, no concurrent load), we can run that when things quiet down. The data above confirms correctness on 122B with thinking mode.

Clean 15-turn benchmark (solo, no concurrent load). M2 Ultra 128GB, Qwen3.5-122B-A10B 5-bit, thinking enabled, vllm-mlx BatchedEngine:

TTFT growth factor: 0.68x. TTFT plateaus around 26s for heavy turns (1500+ chars content) and drops to ~5s for short responses. No degradation across turns.

This is through vllm-mlx's prefix cache (not mlx-lm's LRUPromptCache), so the checkpoint mechanism in this PR isn't directly exercised. For a proper A/B with mlx-lm's server, I'll need to temporarily swap servers. Can do that in a follow-up session.

The cache.py changes (rewind/trim) are confirmed working on 122B MTP with thinking enabled. No issues.

PR #1042 A/B Test: mlx-lm server on 122B (as requested)

Here's the proper A/B using mlx-lm's server (not vllm-mlx), exercising the LRUPromptCache checkpoint logic directly.

Setup: M2 Ultra 128GB, Qwen3.5-122B-A10B 5-bit, python -m mlx_lm.server, thinking enabled, 15-turn coding conversation.

TTFT growth factor: 0.28x (sublinear -- prefix cache is working across turns).

The key signal: Turn 1 cold-starts at 36.6s, Turn 5 drops to 3.7s (10x improvement from cache hit), and by Turn 15 TTFT is only 10.4s despite accumulating ~60K chars of conversation context. Without PR #1042, thinking models would get 0% cache reuse and TTFT would grow linearly with conversation length.

Bug found and fixed

During testing, _can_rewind_layer_cache crashed with a NameError on the legacy fallback path (line 347):

# Line 347 - `rewind` was never defined in scope
if not callable(is_trimmable) or (not callable(trim) and not callable(rewind)):

Fix: add rewind = getattr(layer_cache, "rewind", None) before line 347. This path is hit when a layer cache doesn't implement can_rewind() but does implement the legacy is_trimmable()/trim()/rewind() contract. With Qwen3.5's MTP cache layers, the can_rewind path is used, but the fallback crashes for any third-party cache that only implements the legacy interface.

Code audit: additional findings

While testing on 122B, I audited the prompt cache and batch paths more carefully. Beyond the NameError fix reported above, here are findings worth addressing:

1. _rewind_layer_cache disagrees with _can_rewind_layer_cache (logic bug)

For cache types with is_trimmable()/trim() but no custom can_rewind/rewind (KVCache, RotatingKVCache, etc.):

• _can_rewind_layer_cache correctly falls to the legacy path (line 345+) because can_rewind is None. Returns True.
• _rewind_layer_cache finds rewind is callable (inherited from _BaseCache, raises NotImplementedError). Enters try at line 374, calls rewind(), catches NotImplementedError, returns False. Never reaches its own legacy trim fallback (lines 386-395).

Impact: In fetch_nearest_cache, for any "longer" cached entry with these cache types, _can_rewind says "yes", a copy.deepcopy is made (expensive!), then _rewind returns False, and the deepcopy is discarded. The longer-match rewind optimization is non-functional for all cache types except BatchRotatingKVCache.

Fix: Check _has_rewind_impl() before calling rewind() in _rewind_layer_cache, mirroring the logic in _can_rewind_layer_cache.

2. args.top_logprobs uses wrong request in batch loop (line 1128)

_format_top_logprobs(r.logprobs, args.top_logprobs, current_tokenizer)

args is from the last request dequeued (line 979), not the request that produced batch response r. Each request in the batch could have a different top_logprobs value. Should store top_logprobs per-request when inserted into the batch.

3. CacheOrder.pop() crash on empty deques (line 196-199)

When both _lru and _lru_checkpoints are empty, 0 >= 0 is True, and _lru.popleft() raises IndexError. Called from trim_to which has no guard against empty cache. Trigger: trim_to(n_bytes=0) when _n_bytes is out of sync.

4. Legacy fallback in _rewind_layer_cache is dead code (lines 386-395)

Related to #1. The fallback path is unreachable for any _BaseCache subclass because _BaseCache defines rewind() (raises NotImplementedError), making getattr(layer_cache, "rewind", None) always callable. The if callable(rewind): at line 373 always enters the try/except. Only runs for non-_BaseCache custom caches.

These were found during production testing on Qwen3.5-122B. Happy to discuss any of them.

Thanks for both the isolated mlx_lm.server A/B and the follow-up audit. That 0.28x TTFT growth factor on 122B is exactly the kind of signal that helps de-risk this PR, and I really appreciate you taking the time to swap servers and isolate the run.

On the audit findings:

1 & 4 (_rewind_layer_cache / _can_rewind_layer_cache disagreement + dead legacy fallback): Good catches. These are two faces of the same bug. _rewind_layer_cache now checks _has_rewind_impl() before attempting rewind(), so _BaseCache subclasses with only trim() (KVCache, RotatingKVCache) correctly fall through to the legacy path instead of hitting the base stub and wasting a deepcopy. The missing rewind = getattr(...) on the legacy fallback path is also fixed in the same commit. Six agreement tests were added covering legacy non-_BaseCache caches, _BaseCache-with-trim()-but-no-rewind, and BatchRotatingKVCache. Pushed in aa0310b.

2 (top_logprobs batch scoping): Agreed — pre-existing and worth a separate fix.

3 (CacheOrder.pop() on empty deques): Opened separately as #1054 off upstream/main, with no dependency on this PR.

And thanks again for digging deeply enough to catch the legacy fallback NameError.

@angeloskath @awni — this PR has been open since March 18 with no maintainer review. The author (lyonsno) has been very responsive to feedback, addressing all code review findings promptly including bug fixes with new test coverage.

Is there a concern with the approach or scope? Happy to help if anything is needed to move this forward.

Context: thinking models currently get 0% prompt cache reuse on multi-turn conversations because thinking tokens diverge between turns. This PR fixes that with checkpoint-based rewind. We validated on 122B production (0.28x TTFT growth factor across 15 turns) and found + reported 5 bugs during testing, all addressed by the author.

Sorry for taking long to reply. #1072 should be taking care of this. Big part of it was already taken care of on main but either way feel free to provide a snippet on #1072 that shows the cache not being properly used if you are still experiencing issues.